import os, gzip, pickle
from pathlib import Path
from typing import Iterable, Type, Optional, Union, Literal
from feasytools import TableWriter, ReadOnlyTable
import numpy as np
from fplugins import *
from .logcs import *
from .logev import *
from .loggr import *

FILE_TL = "time_line.pkl.gz"

def _suffixchk(suf):
    if suf is None: return False
    assert suf in TableFileSuffixList
    return True

TableFileSuffix=Literal['<mem>','.sdt','.sdt.gz','.csv']
TableFileSuffixList=['<mem>','.sdt','.sdt.gz','.csv']

class StaMan:
    __ava_logger:dict[str,Type] = {
        "cspk": StaCSPK,
        "ev": StaEV,
        "gen": StaGen,
        "bus": StaBus,
    }
    @staticmethod
    def Register(name:str,base:Type)->None:
        '''注册一个统计项'''
        if name in StaMan.__ava_logger: raise ValueError(f"已经注册过统计项{name}")
        assert issubclass(base, StaBase)
        StaMan.__ava_logger[name] = base
    
    @staticmethod
    def Get(name:str)->Type:
        '''获取一个统计项'''
        return StaMan.__ava_logger[name]

class StaWriter:
    '''统计数据记录器'''
    __items:dict[str,StaBase]
    __time:list[int]
    
    @property
    def TimeLine(self)->list[int]: return self.__time

    def Writer(self,name:str)->TableWriter: return self.__items[name].Writer

    def __init__(self,path:str,suffix:str,tinst:TrafficInst,plugins:dict[str,PluginBase]):
        self.__path = path
        _suffixchk(suffix)
        self.__suffix = suffix
        self.__time = []
        self.__items = {}
        self.__inst = tinst
        self.__plug = plugins
    
    def Add(self,sta_name:str)->None:
        '''添加一个统计项, 从StaMan的注册项中选择'''
        sta_type = StaMan.Get(sta_name)
        if sta_name in self.__items: raise ValueError("已经注册过该项")
        self.__items[sta_name] = sta_type(self.__path,self.__suffix,self.__inst,self.__plug)

    def Log(self,time:int):
        self.__time.append(time)
        for item in self.__items.values():
            try:
                item.LogOnce()
            except Exception as e:
                print(f"记录{item._name}时发生错误: ")
                raise e
    
    def close(self):
        with gzip.open(str(Path(self.__path)/FILE_TL),"wb") as fp:
            pickle.dump(self.__time,fp)
        for item in self.__items.values():
            try:
                item.close()
            except Exception as e:
                print(f"关闭{item._name}时发生错误: ")
                raise e

class StaReader:
    f'''**只读**数据统计模块, 只能从文件夹创建'''
    def __init__(self,path:str,time_line:Optional[list[int]]=None):
        '''
        初始化
            path: 初始化文件夹
            time_line: 时间点列表
        '''
        work_dir = Path(path)
        dir_con = os.listdir(path)
        self.__items:dict[str,ReadOnlyTable] = {}
        for file in dir_con:
            if file.lower() == FILE_TL and time_line is None:
                with gzip.open(str(work_dir/file)) as fp:
                    time_line = pickle.load(fp)
                    assert isinstance(time_line,list)
            elif file.endswith('.sdt') or file.endswith('.sdt.gz') or file.endswith('.csv'):
                fname = file.removesuffix('.sdt').removesuffix('.sdt.gz').removesuffix('.csv')
                if StaMan.Get(fname) is not None:
                    self.__items[fname] = ReadOnlyTable(str(work_dir/file),np.float32)
            else:
                continue
        if time_line is None: raise FileNotFoundError("未找到时间线文件")
        self.__time = time_line
        
    @property
    def TimeLine(self)->list[int]: 
        return self.__time
    
    def __contains__(self,table_name:str)->bool:
        return table_name in self.__items
    
    def Headof(self,table_name)->list[str]:
        '''获取一个统计表的列名'''
        return self.__items[table_name].head
    
    def GetTable(self,tabel_name:str)->ReadOnlyTable:
        '''获取一个统计表'''
        return self.__items[tabel_name]
    
    def GetColumn(self,tabel_name:str,cols:Union[str,int])->np.ndarray:
        '''获取一个统计表的列'''
        return self.__items[tabel_name].col(cols)
    
    def GetColumns(self,tabel_name:str,cols:Iterable[Union[str,int]])->np.ndarray:
        '''获取一个统计表的多列'''
        return self.__items[tabel_name].col(cols)
    
    def GetColumnNames(self,tabel_name:str)->list[str]:
        '''获取一个统计表的列名'''
        return self.__items[tabel_name].head
    
    def ForceLoad(self,table_name:str):
        '''强制加载一个表'''
        self.__items[table_name].data